import numpy as np

# 默认训练集
default_learning_dataset = {
    "宝贝当家": [45, 2, 9, "喜剧片"],
    "美人鱼": [21, 17, 5, "喜剧片"],
    "澳门风云3": [54, 9, 11, "喜剧片"],
    "功夫熊猫3": [39, 0, 31, "喜剧片"],
    "谍影重重": [5, 2, 57, "动作片"],
    "叶问3": [3, 2, 65, "动作片"],
    "伦敦陷落": [2, 3, 55, "动作片"],
    "我的特工爷爷": [6, 4, 21, "动作片"],
    "奔爱": [7, 46, 4, "爱情片"],
    "夜孔雀": [9, 39, 8, "爱情片"],
    "代理情人": [9, 38, 2, "爱情片"],
    "新步步惊心": [8, 34, 17, "爱情片"]
}


def euclidean_distance(point_a, point_b):
    """
    计算两点之间的欧式距离

    :param point_a: A 点坐标
    :param point_b: B 点坐标
    :return: 两点之间的欧式距离
    """
    if len(point_a) != len(point_b):
        raise ValueError("两点维度不一致")
    return np.sqrt(np.sum(np.power(np.array(point_a) - np.array(point_b), 2)))


def label_count(label_list):
    """
    标签计数

    :param label_list: 标签列表
    :return: 标签计数字典 { 'label': count }
    """
    count_dict = {}
    for label in label_list:
        if label not in count_dict.keys():
            count_dict[f"{label}"] = 1
        else:
            count_dict[f"{label}"] += 1
    return count_dict


def predict_movie_genre_with_knn_classifier(move_dataset, learning_dataset=None, k=6):
    """
    通过 KNN 分类器预测电影类别

    :param move_dataset: 电影集
    :param learning_dataset: 训练集
    :param k: 最大近邻数 (default: 6)
    :return: 预测结果, 更新电影类别后的电影集
    """
    if learning_dataset is None:
        learning_dataset = default_learning_dataset

    predict_result = {}
    updated_move_dataset = move_dataset.copy()
    for move_name, move_data in move_dataset.items():
        distance_list = []
        for learning_move_name, learning_move_data in learning_dataset.items():
            distance_list.append((euclidean_distance(move_data[:3], learning_move_data[:3]), learning_move_data[-1]))
        distance_list.sort(key=lambda distance: distance[0])
        possible_label_list = [label for _, label in distance_list[:k]]
        label_count_dict = label_count(possible_label_list)
        predict_label = max(label_count_dict, key=label_count_dict.get)
        predict_result[move_name] = predict_label
        updated_move_dataset[move_name][-1] = predict_label
    return predict_result, updated_move_dataset


if __name__ == '__main__':
    # 测试集
    testing_dataset = {
        "老友记": [29, 10, 2, ""]
    }
    print(predict_movie_genre_with_knn_classifier(testing_dataset)[0])
    print(predict_movie_genre_with_knn_classifier(testing_dataset)[1])
